import os
import re
import shutil
import zipfile
from concurrent.futures import ThreadPoolExecutor
try:
    import gdal
except:
    from osgeo import gdal

os.environ['PROJ_LIB'] = r"/home/冯洋洋/PythonEnvironment/RasterShapeEnv/share/proj"

# unzip zip to file
def get_l1c_png(zip_path, save_dir):
    zip_name = os.path.basename(zip_path)
    tile = zip_name.split("_")[5][1:]
    tile_dir = os.path.join(save_dir, tile)
    if not os.path.exists(tile_dir):
        os.makedirs(tile_dir)
    update_dir = os.path.join(tile_dir, "update")
    if not os.path.exists(update_dir):
        os.makedirs(update_dir)        
    namere = "T[\w]{5}_[\d]{8}T[\d]{6}_PVI\.jp2"
    zipFile = zipfile.ZipFile(zip_path, "r")
    for fileInfo in zipFile.infolist():
        if re.match(namere, os.path.basename(fileInfo.filename)):
            jp2_path = zipFile.extract(fileInfo.filename, tile_dir)
            zipFile.close()
            png_path = os.path.join(tile_dir, zip_name.replace(".zip", ".png"))
            gdal.Translate(png_path, jp2_path, format='PNG')
            shutil.rmtree(os.path.join(tile_dir, zip_name.replace(".zip", '.SAFE')))
            os.remove(os.path.join(tile_dir, zip_name.replace(".zip", '.png.aux.xml')))
            break      
    return 

def main():
    # zip_dir = input("请输入Sentinel-2号level-1产品zip目录:")
    # img_dir = input("请输入Sentinel-2号level-1产品img目录:")
    # for zip_name in os.listdir(zip_dir):
    #     zip_path = os.path.join(zip_dir, zip_name)
    #     get_l1c_png(zip_path, img_dir)
    zip_dir = r"/83To82/Yang/data/zip"
    img_dir = r"/home/遥感技术部/遥感数据/冯洋洋/png"
    zip_paths = [
        os.path.join(zip_dir, zip_name)
        for zip_name in os.listdir(zip_dir)
    ]
    with ThreadPoolExecutor(10) as executor:
        executor.map(get_l1c_png, zip_paths, [img_dir]*len(zip_paths))
    return


if __name__ == "__main__":
    main()

